"""
This class represents a simple neural network. For now, it's just a list
of layers.
"""
import numpy as np

class NeuralNetwork:

    def __init__(self):
        self.layers = []

    def add_layer(self, layer):
        self.layers.append(layer)

    def forward_pass(self, batch_data):
        """
        Perform a forward pass through the network by passing the data through
        each layer.
        """
        input_data = batch_data
        for layer in self.layers:
            input_data = layer.forward_pass(input_data)
        return input_data
    
    def backward_pass(self, dvalues):
        """
        Perform a backward pass through the network given gradient at output.
        """
        grad = dvalues
        for layer in reversed(self.layers):
            grad = layer.backward_pass(grad)

    def train(self, X, y, loss_obj, epochs, learning_rate, batch_size=None, shuffle=True,
              verbose=False, X_val=None, y_val=None, metric=None, target_std=None, plot=False):
        """
        Train the network using gradient descent.
        X: input data of shape (features, samples)
        y: true labels of shape (outputs, samples)
        loss_obj: instance of a loss class with forward_pass and backward_pass
        epochs: number of epochs to train
        learning_rate: step size for parameter updates
        batch_size: number of samples per batch (None for full-batch)
        shuffle: whether to shuffle data each epoch
        verbose: whether to print progress
        X_val: validation data of shape (features, samples), optional
        y_val: validation labels of shape (outputs, samples), optional
        metric: string, either "accuracy" or "rmse", to compute on validation data, optional
        target_std: per-output std used to standardize y, shape (outputs, 1), optional.
            If provided with metric="rmse", the printed RMSE is reported in the
            original target units instead of standardized units.
        plot: if True, pop up a matplotlib figure with training diagnostics
            (train loss, val loss, val metric, weight L2 norm) when training
            finishes or is interrupted (default False).

        Training can be interrupted at any time with Ctrl+C, and the method will return gracefully.
        """
        n_samples = X.shape[1]
        train_losses = []
        val_losses = []
        val_metrics = []
        weight_norms = []
        metric_name = None

        try:
            for epoch in range(1, epochs + 1):
                # Shuffle data
                if shuffle:
                    perm = np.random.permutation(n_samples)
                    X_shuffled = X[:, perm]
                    y_shuffled = y[:, perm]
                else:
                    X_shuffled = X
                    y_shuffled = y

                # Determine batch size
                bs = batch_size or n_samples
                n_batches = int(np.ceil(n_samples / bs))
                epoch_loss = 0.0

                for batch_idx in range(n_batches):
                    start = batch_idx * bs
                    end = start + bs
                    X_batch = X_shuffled[:, start:end]
                    y_batch = y_shuffled[:, start:end]

                    # Forward pass
                    y_pred = self.forward_pass(X_batch)
                    loss = loss_obj.forward_pass(y_pred, y_batch)
                    epoch_loss += loss * X_batch.shape[1]

                    # Backward pass
                    dvalues = loss_obj.backward_pass()
                    self.backward_pass(dvalues)

                    # Update parameters
                    for layer in self.layers:
                        if hasattr(layer, 'weights'):
                            layer.weights -= learning_rate * layer.dweights
                            layer.biases  -= learning_rate * layer.dbiases

                # Average loss over all samples
                epoch_loss /= n_samples
                train_losses.append(epoch_loss)

                # Track total weight L2 norm so we can plot stability over time.
                sum_sq = 0.0
                for layer in self.layers:
                    if hasattr(layer, 'weights'):
                        sum_sq += np.sum(layer.weights ** 2)
                weight_norms.append(float(np.sqrt(sum_sq)))

                # Compute validation metric if validation data and metric are provided
                # this is just so we can watch the training quality as we run
                if X_val is not None and y_val is not None and metric is not None:
                    y_pred_val = self.forward_pass(X_val)
                    val_loss = loss_obj.forward_pass(y_pred_val, y_val)
                    val_losses.append(val_loss)

                    if metric == "accuracy":
                        # Compute accuracy for classification
                        # Assumes y_val is one-hot encoded or integer labels
                        if y_val.shape[0] > 1:  # One-hot encoded
                            y_pred_labels = np.argmax(y_pred_val, axis=0)
                            y_true_labels = np.argmax(y_val, axis=0)
                        else:  # Integer labels
                            y_pred_labels = np.round(y_pred_val).astype(int)
                            y_true_labels = y_val.astype(int)
                        
                        val_metric = np.mean(y_pred_labels == y_true_labels)
                        metric_name = "accuracy"
                    
                    elif metric == "rmse":
                        # Compute RMSE for regression. If targets were standardized
                        # by dividing by target_std, scale the standardized errors
                        # back so the reported RMSE is in original target units.
                        if target_std is not None:
                            sq_err_real = ((y_pred_val - y_val) * target_std) ** 2
                            val_metric = np.sqrt(np.mean(sq_err_real))
                        else:
                            val_metric = np.sqrt(val_loss)
                        metric_name = "rmse"
                    
                    else:
                        raise ValueError(f"Unsupported metric: {metric}. Use 'accuracy' or 'rmse'.")

                    val_metrics.append(val_metric)

                    if verbose:
                        print(f"Epoch {epoch}/{epochs} - loss: {epoch_loss:.8f} - val_loss: {val_loss:.8f} - val_{metric_name}: {val_metric:.8f}")

                elif verbose:
                    print(f"Epoch {epoch}/{epochs} - loss: {epoch_loss:.8f}")

        except KeyboardInterrupt:
            if verbose:
                print("\nTraining interrupted by user. Returning with current model state.")

        if plot:
            self._plot_training(train_losses, val_losses, val_metrics, weight_norms, metric_name)

    def _plot_training(self, train_losses, val_losses, val_metrics, weight_norms, metric_name):
        """
        Pop up a 2x2 matplotlib figure summarizing training:
        train loss, validation loss, validation metric, and total weight L2 norm.
        """
        import matplotlib.pyplot as plt

        fig, axes = plt.subplots(2, 2, figsize=(12, 8))

        # Top-left: train loss per epoch
        axes[0, 0].plot(range(1, len(train_losses) + 1), train_losses, color='C0')
        axes[0, 0].set_title('Train Loss')
        axes[0, 0].set_xlabel('Epoch')
        axes[0, 0].set_ylabel('Loss')
        axes[0, 0].grid(True, alpha=0.3)

        # Top-right: validation loss per epoch
        if val_losses:
            axes[0, 1].plot(range(1, len(val_losses) + 1), val_losses, color='C1')
            axes[0, 1].set_title('Validation Loss')
        else:
            axes[0, 1].set_title('Validation Loss (no validation data)')
        axes[0, 1].set_xlabel('Epoch')
        axes[0, 1].set_ylabel('Loss')
        axes[0, 1].grid(True, alpha=0.3)

        # Bottom-left: validation metric (accuracy or rmse) per epoch
        if val_metrics and metric_name:
            axes[1, 0].plot(range(1, len(val_metrics) + 1), val_metrics, color='C2')
            axes[1, 0].set_title(f'Validation {metric_name.capitalize()}')
            axes[1, 0].set_ylabel(metric_name)
        else:
            axes[1, 0].set_title('Validation Metric (none)')
            axes[1, 0].set_ylabel('value')
        axes[1, 0].set_xlabel('Epoch')
        axes[1, 0].grid(True, alpha=0.3)

        # Bottom-right: total weight L2 norm — shows whether weights stay
        # well-bounded or drift toward instability.
        axes[1, 1].plot(range(1, len(weight_norms) + 1), weight_norms, color='C3')
        axes[1, 1].set_title('Total Weight L2 Norm')
        axes[1, 1].set_xlabel('Epoch')
        axes[1, 1].set_ylabel('||W||')
        axes[1, 1].grid(True, alpha=0.3)

        fig.tight_layout()
        plt.show()